#include "CNN.h"
#include <algorithm>
#include <cmath>

float X_scaler_mean[8];
float X_scaler_scale[8];
float y_scaler_mean;
float y_scaler_scale;
std::vector<std::vector<float>> input;

void init_variables() {
    input = {
        {0.203, 0.203, 0.203, 0.203, 0.202, 0.201, 0.201, 0.198, 0.191, 0.184},  
        {0.095, 0.095, 0.095, 0.095, 0.095, 0.094, 0.094, 0.093, 0.091, 0.088},  
        {0.008, 0.008, 0.008, 0.008, 0.008, 0.008, 0.008, 0.007, 0.007, 0.007},  
        {0.078, 0.078, 0.078, 0.078, 0.078, 0.078, 0.078, 0.078, 0.078, 0.077},  
        {0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001},  
        {0.017, 0.017, 0.017, 0.017, 0.017, 0.017, 0.017, 0.017, 0.017, 0.017},  
        {0.004, 0.004, 0.004, 0.004, 0.004, 0.004, 0.004, 0.004, 0.004, 0.004}, 
        {0.032, 0.032, 0.032, 0.032, 0.032, 0.032, 0.032, 0.032, 0.032, 0.031}
    };

    X_scaler_mean[0] = 0.6682482944139136;
    X_scaler_mean[1] = 0.41993549107142764;
    X_scaler_mean[2] = 0.28496372481683674;
    X_scaler_mean[3] = 0.4328882612179549;
    X_scaler_mean[4] = 0.47025940934063604;
    X_scaler_mean[5] = 0.6220351247710388;
    X_scaler_mean[6] = 0.3759382497710415;
    X_scaler_mean[7] = 0.6854925538003623;

    X_scaler_scale[0] = 0.5693549442920608;
    X_scaler_scale[1] = 0.5557024917321426;
    X_scaler_scale[2] = 0.24714607362523378;
    X_scaler_scale[3] = 0.3225593908236477;
    X_scaler_scale[4] = 0.7283238561974298;
    X_scaler_scale[5] = 0.14123522588459905;
    X_scaler_scale[6] = 0.5750555055965144;
    X_scaler_scale[7] = 0.7091560850845181;

    y_scaler_mean = 90.98443223443223;
    y_scaler_scale = 209.32142619570598;
}

MultiTaskModel::MultiTaskModel() {
    initWeights();  
}

void MultiTaskModel::initWeights() {
    conv1_weights = {
        {
            {
                -0.16104619204998016,
                0.14749681949615479,
                0.13192063570022583
            },
            {
                0.250449538230896,
                0.4528956711292267,
                -0.7018706202507019
            },
            {
                0.34893837571144104,
                0.2454437017440796,
                0.5517166256904602
            },
            {
                0.2136150598526001,
                0.429021954536438,
                -0.30771273374557495
            },
            {
                -0.288152277469635,
                -0.4148397147655487,
                -0.5754204988479614
            },
            {
                0.026125211268663406,
                -0.09746038168668747,
                -0.2579357922077179
            },
            {
                0.11645495891571045,
                0.38625356554985046,
                -0.0600283220410347
            },
            {
                -0.45624276995658875,
                -0.2152085155248642,
                0.06621681898832321
            }
        },
        {
            {
                -0.48990729451179504,
                -0.6946542263031006,
                0.6451901197433472
            },
            {
                0.1826496124267578,
                0.29014459252357483,
                0.5777813196182251
            },
            {
                -0.0011158575071021914,
                0.2674388885498047,
                0.3357655107975006
            },
            {
                0.1993824988603592,
                0.13104605674743652,
                0.46075308322906494
            },
            {
                0.07424528896808624,
                0.03578152507543564,
                0.2591502368450165
            },
            {
                0.5167933702468872,
                0.3585835099220276,
                0.4739592671394348
            },
            {
                0.4405916929244995,
                0.36316734552383423,
                0.5503840446472168
            },
            {
                -0.3184254467487335,
                -0.721964418888092,
                -0.3280949294567108
            }
        },
        {
            {
                0.21403169631958008,
                0.649116575717926,
                0.5417736172676086
            },
            {
                0.10763484984636307,
                0.13426463305950165,
                0.11990568786859512
            },
            {
                0.3066021800041199,
                -0.04832958057522774,
                0.06713550537824631
            },
            {
                0.18521088361740112,
                0.40583765506744385,
                0.29169148206710815
            },
            {
                -0.10555856674909592,
                -0.18564777076244354,
                0.03351583331823349
            },
            {
                0.19007933139801025,
                0.18539798259735107,
                0.21486373245716095
            },
            {
                0.11615370213985443,
                -0.02905130386352539,
                0.08033738285303116
            },
            {
                -0.22186899185180664,
                -0.27649688720703125,
                -0.25963494181632996
            }
        },
        {
            {
                0.21403881907463074,
                -0.4985678791999817,
                -1.2674452066421509
            },
            {
                0.3597254753112793,
                -0.19742752611637115,
                -0.2981821894645691
            },
            {
                0.6933437585830688,
                0.34388795495033264,
                0.7114121317863464
            },
            {
                -0.10834582895040512,
                -0.18372218310832977,
                -0.2310866415500641
            },
            {
                -0.42884644865989685,
                -0.0880763977766037,
                0.02015446312725544
            },
            {
                -0.4427637755870819,
                -0.19690538942813873,
                0.3320031762123108
            },
            {
                -0.15006950497627258,
                0.006054629106074572,
                -1.1783157587051392
            },
            {
                0.17095376551151276,
                0.04006732255220413,
                0.5366925597190857
            }
        },
        {
            {
                0.4497618079185486,
                -0.4349999725818634,
                -0.9534463286399841
            },
            {
                0.3188202977180481,
                -0.6872991323471069,
                -0.01235381979495287
            },
            {
                -0.796009361743927,
                0.09253213554620743,
                0.12996971607208252
            },
            {
                0.07318034023046494,
                0.3391549289226532,
                -0.4657996892929077
            },
            {
                -0.23928503692150116,
                0.37803542613983154,
                0.45338040590286255
            },
            {
                0.48308730125427246,
                -0.3093600273132324,
                -0.32966554164886475
            },
            {
                -0.8632999062538147,
                -0.7849550247192383,
                -1.2005763053894043
            },
            {
                -0.23835459351539612,
                0.19672705233097076,
                0.07099705934524536
            }
        },
        {
            {
                -0.03398504480719566,
                -0.44699713587760925,
                -0.5830332040786743
            },
            {
                0.15212063491344452,
                -0.42270001769065857,
                -1.3757574558258057
            },
            {
                0.6101176738739014,
                0.26118040084838867,
                -0.2666397988796234
            },
            {
                0.5381285548210144,
                -0.43985438346862793,
                -0.1201900914311409
            },
            {
                0.01304006576538086,
                -0.272029310464859,
                0.35378873348236084
            },
            {
                0.11836793273687363,
                -0.21822930872440338,
                -0.4305398762226105
            },
            {
                0.10745050758123398,
                -0.5127661228179932,
                -1.1568262577056885
            },
            {
                0.030573345720767975,
                0.19381767511367798,
                0.7577446103096008
            }
        },
        {
            {
                0.2877923846244812,
                -0.5920714735984802,
                -1.0945580005645752
            },
            {
                -0.10708208382129669,
                -0.4459197223186493,
                -0.3217442035675049
            },
            {
                0.5490514039993286,
                0.0603230744600296,
                -0.17077681422233582
            },
            {
                0.474538654088974,
                -0.08857059478759766,
                -0.69021075963974
            },
            {
                0.1789722442626953,
                -0.019225427880883217,
                0.35351964831352234
            },
            {
                -0.284916490316391,
                -0.47727319598197937,
                0.29924845695495605
            },
            {
                0.3149968087673187,
                0.39389854669570923,
                0.020506571978330612
            },
            {
                -0.019942468032240868,
                0.027432655915617943,
                0.9770858287811279
            }
        },
        {
            {
                0.008434879593551159,
                0.24791952967643738,
                0.016517600044608116
            },
            {
                -0.13383588194847107,
                -0.15069648623466492,
                0.274243026971817
            },
            {
                0.16571560502052307,
                -0.13546693325042725,
                -0.013661980628967285
            },
            {
                0.26855072379112244,
                0.13454976677894592,
                -0.2953815460205078
            },
            {
                0.05967212840914726,
                0.10740657895803452,
                0.059568699449300766
            },
            {
                -0.12892135977745056,
                -0.18514548242092133,
                -0.029932815581560135
            },
            {
                0.13307221233844757,
                0.16855749487876892,
                0.43483009934425354
            },
            {
                0.5004286170005798,
                0.6094614863395691,
                0.44404733180999756
            }
        }
    };
    conv1_bias = {
        -0.365825891494751,
        -1.255890130996704,
        -0.17185984551906586,
        0.47847095131874084,
        -0.15782947838306427,
        0.2196320742368698,
        0.5604979395866394,
        0.5668585896492004
    };

    shared_weights = {
        {
            -0.6240329742431641,
            -1.8725870847702026,
            -2.58512282371521,
            -3.3486781120300293,
            0.27795225381851196,
            0.050046321004629135,
            0.15015092492103577,
            0.152760848402977,
            -0.08879773318767548,
            0.015680819749832153,
            0.13183791935443878,
            0.2633543610572815,
            0.030311422422528267,
            -0.07893719524145126,
            0.13389961421489716,
            -0.258004367351532,
            1.1129157543182373,
            0.912146806716919,
            0.11838909238576889,
            -2.326526165008545,
            -1.5652798414230347,
            0.5065456032752991,
            0.6003890037536621,
            -0.36439791321754456,
            -0.30668801069259644,
            -0.5891534090042114,
            -0.11962166428565979,
            0.4441426396369934,
            0.1901271790266037,
            -0.286299467086792,
            0.015381076373159885,
            0.6864927411079407,
            0.6700565218925476,
            0.5867089033126831,
            -0.6855913996696472,
            -0.19269903004169464,
            0.06107241287827492,
            -0.12065226584672928,
            -0.07535707950592041,
            -0.009833917021751404
        },
        {
            0.6833029389381409,
            0.44532662630081177,
            0.32966890931129456,
            -0.12989726662635803,
            0.8935887813568115,
            0.12878090143203735,
            0.12835368514060974,
            0.23888297379016876,
            0.4580843448638916,
            0.6331474184989929,
            -0.2493424117565155,
            -0.2643760144710541,
            0.008911555632948875,
            -0.08148583024740219,
            0.214008629322052,
            0.2814767062664032,
            0.3619561195373535,
            0.1807508021593094,
            0.5079973340034485,
            0.1151021420955658,
            0.1261388659477234,
            0.17407087981700897,
            -0.019370712339878082,
            0.02241414412856102,
            0.19742600619792938,
            0.11539752781391144,
            0.09177078306674957,
            -0.009160813875496387,
            0.2872017025947571,
            0.5875236392021179,
            0.054674968123435974,
            -0.11143123358488083,
            0.03756733238697052,
            0.663332462310791,
            -0.11787469685077667,
            -0.38733822107315063,
            -0.20078100264072418,
            -0.1908552348613739,
            0.04699741676449776,
            0.21023041009902954
        },
        {
            0.08700184524059296,
            -0.28639325499534607,
            -0.7838150262832642,
            -1.1179579496383667,
            -0.5546892285346985,
            0.6916200518608093,
            0.3893960416316986,
            0.10257033258676529,
            -0.4815126061439514,
            -1.0446844100952148,
            0.5976892709732056,
            0.30269449949264526,
            0.08863646537065506,
            -0.06764993071556091,
            -0.992875874042511,
            0.09733632206916809,
            0.10966214537620544,
            0.15517748892307281,
            0.2067287117242813,
            0.25646692514419556,
            0.022993478924036026,
            0.263287216424942,
            -0.0668465867638588,
            -0.02663537859916687,
            0.4774092435836792,
            -0.046136002987623215,
            0.10222258418798447,
            0.0976908877491951,
            -0.18057793378829956,
            0.6412362456321716,
            0.12571191787719727,
            -0.07213153690099716,
            0.03544718399643898,
            0.3831160366535187,
            0.4916296601295471,
            0.2037772834300995,
            -0.01186726801097393,
            0.026156339794397354,
            0.1461368352174759,
            0.1053302064538002
        },
        {
            -0.874483048915863,
            -1.4110238552093506,
            -1.0462080240249634,
            0.7868564128875732,
            -0.5266740918159485,
            0.21384161710739136,
            0.25811585783958435,
            0.033040568232536316,
            -0.16963812708854675,
            -0.20548395812511444,
            0.23174035549163818,
            -0.004807057324796915,
            0.07752019166946411,
            0.15986406803131104,
            0.03287563472986221,
            0.18377262353897095,
            -0.19275794923305511,
            -0.23365558683872223,
            -0.4364017844200134,
            -0.29900848865509033,
            -0.3137190639972687,
            -0.039639052003622055,
            0.2783242166042328,
            -0.5351633429527283,
            -1.1731693744659424,
            -0.08640400320291519,
            0.03576989099383354,
            -0.1425144076347351,
            0.4632021486759186,
            2.1182923316955566,
            0.298117458820343,
            0.5953034162521362,
            0.6270142197608948,
            -2.2101995944976807,
            -0.06117376685142517,
            0.22238744795322418,
            -0.008733605034649372,
            -0.07523402571678162,
            -0.2282850444316864,
            -0.19087526202201843
        },
        {
            -0.6834465861320496,
            -0.11598792672157288,
            0.03173871338367462,
            0.44698598980903625,
            -0.5958941578865051,
            0.2643401324748993,
            0.35145705938339233,
            0.1897319108247757,
            -0.05099176988005638,
            -0.11974524706602097,
            0.0731365904211998,
            0.07176000624895096,
            0.22781382501125336,
            0.27071860432624817,
            0.26876574754714966,
            0.024554714560508728,
            0.22068987786769867,
            -0.11594787240028381,
            0.2807988226413727,
            1.0249360799789429,
            0.09794801473617554,
            0.32149821519851685,
            -0.007692273240536451,
            0.0660257637500763,
            -0.027705227956175804,
            -0.42276811599731445,
            -0.12274942547082901,
            -0.31181222200393677,
            0.096322201192379,
            0.22792930901050568,
            0.26630789041519165,
            -0.07015452533960342,
            -0.2712489664554596,
            -0.11464696377515793,
            0.7802014350891113,
            0.13778255879878998,
            0.023486895486712456,
            0.1345032900571823,
            -0.058348942548036575,
            -0.2852822244167328
        },
        {
            0.2821296155452728,
            -0.047498371452093124,
            -0.2523207664489746,
            -0.9268863797187805,
            0.3329358994960785,
            0.004681237041950226,
            -0.03803045302629471,
            0.051603324711322784,
            0.30153656005859375,
            0.2966310977935791,
            -0.0678834319114685,
            -0.18393032252788544,
            -0.1492283195257187,
            0.04409154877066612,
            -0.06291329860687256,
            0.11550132930278778,
            -0.010053045116364956,
            -0.04586653411388397,
            -0.4238894283771515,
            -0.970139741897583,
            0.011544546112418175,
            -0.1759706735610962,
            -0.4391252100467682,
            -0.9833240509033203,
            0.568818986415863,
            0.7525056600570679,
            0.7101227045059204,
            0.660209059715271,
            -0.07910233736038208,
            -0.3498871624469757,
            0.045480843633413315,
            0.24051308631896973,
            0.033177584409713745,
            0.10044028609991074,
            -0.910814106464386,
            -0.029280854389071465,
            0.021872559562325478,
            0.2473633736371994,
            0.2525567412376404,
            0.586700975894928
        },
        {
            -0.7235572934150696,
            -0.1607043743133545,
            -0.06821364909410477,
            0.6913576126098633,
            -0.29720404744148254,
            0.25616276264190674,
            0.29962530732154846,
            0.044240280985832214,
            -0.004135752562433481,
            0.02577897161245346,
            0.33383819460868835,
            0.24215587973594666,
            0.22019970417022705,
            0.42796504497528076,
            0.28725773096084595,
            0.25908640027046204,
            0.24665024876594543,
            0.05302714928984642,
            0.36288467049598694,
            0.8503999710083008,
            0.32677194476127625,
            0.3247361481189728,
            0.06147100776433945,
            -0.0043905009515583515,
            -0.17560997605323792,
            -0.5438954830169678,
            0.005776382051408291,
            -0.2687564492225647,
            0.055049918591976166,
            0.12636065483093262,
            0.08579988777637482,
            -0.09854280948638916,
            -0.4559900760650635,
            -0.3533054292201996,
            0.6916165947914124,
            0.23523752391338348,
            0.11505841463804245,
            -0.06749381124973297,
            -0.18487630784511566,
            -0.2781944274902344
        },
        {
            -0.13888758420944214,
            -0.14946220815181732,
            0.0366663858294487,
            -0.07595758140087128,
            0.10131184756755829,
            0.09656975418329239,
            -0.008640775457024574,
            0.0006701372913084924,
            -0.014770083129405975,
            -0.04596995934844017,
            -0.08706303685903549,
            -0.07027719169855118,
            -0.013953202404081821,
            0.07414418458938599,
            -0.11293207108974457,
            -0.08733218163251877,
            -0.08402608335018158,
            -0.10561968386173248,
            -0.004353920463472605,
            0.05240369960665703,
            0.11541488021612167,
            0.05125043913722038,
            -0.11636099219322205,
            -0.08486047387123108,
            -0.034739427268505096,
            -0.03960256651043892,
            -0.0074236635118722916,
            -0.1451440304517746,
            0.04921211302280426,
            -0.16805683076381683,
            -0.1504538208246231,
            -0.17887310683727264,
            0.08493895083665848,
            -0.18747256696224213,
            0.07642728835344315,
            -0.06743898242712021,
            -0.07935580611228943,
            0.039358023554086685,
            -0.07317876070737839,
            0.0421622171998024
        },
        {
            -0.15378537774085999,
            -0.28730159997940063,
            -0.030032021924853325,
            0.5152372717857361,
            0.18171322345733643,
            0.15189309418201447,
            0.09394199401140213,
            -0.09587323665618896,
            -0.11829843372106552,
            -0.12141498178243637,
            0.07004702836275101,
            -0.05277976021170616,
            -0.22371046245098114,
            -0.08309145271778107,
            -0.273596853017807,
            0.1655302345752716,
            0.05470208823680878,
            0.1571304202079773,
            0.07810550928115845,
            0.3807794451713562,
            0.14317969977855682,
            0.06545119732618332,
            0.4006801247596741,
            0.3726992607116699,
            0.1702776849269867,
            0.20149756968021393,
            0.19166944921016693,
            0.3564024269580841,
            0.2502138614654541,
            0.31206098198890686,
            0.24820169806480408,
            0.22192814946174622,
            0.30562615394592285,
            0.06550010293722153,
            0.8151507377624512,
            0.2972777783870697,
            0.1989353448152542,
            0.22418569028377533,
            0.02464776486158371,
            -0.33212539553642273
        },
        {
            -0.33940836787223816,
            -0.22425724565982819,
            -0.13211195170879364,
            0.5623149275779724,
            -0.64549320936203,
            0.21715572476387024,
            0.20260334014892578,
            0.01060885563492775,
            -0.005728860851377249,
            -0.037907663732767105,
            0.3451240062713623,
            0.1449989676475525,
            0.12207657098770142,
            0.15125957131385803,
            -0.07999899983406067,
            0.16292612254619598,
            0.20069758594036102,
            -0.08541962504386902,
            0.059138476848602295,
            0.4501041769981384,
            0.1284191608428955,
            0.07811524718999863,
            0.02416890114545822,
            0.13245433568954468,
            -0.14363710582256317,
            0.08244768530130386,
            0.17387358844280243,
            0.17515389621257782,
            0.049930356442928314,
            0.1899675875902176,
            0.08350925892591476,
            -0.07333853095769882,
            0.005439543630927801,
            -0.26622655987739563,
            0.6500034928321838,
            0.20782135426998138,
            -0.009149020537734032,
            0.12392596155405045,
            -0.24542756378650665,
            -0.38277995586395264
        }
    };

    shared_bias = {
        0.7769724130630493,
        -0.5532165169715881,
        0.4801519513130188,
        0.4939594268798828,
        0.12540137767791748,
        0.48637256026268005,
        -0.22746694087982178,
        -0.16601121425628662,
        0.27538585662841797,
        0.34242069721221924
    };

    regression_head_weights = {
        {
            0.5330276489257812,
            0.6666262149810791,
            0.7240488529205322,
            -1.6222376823425293,
            -0.24027207493782043,
            0.592673122882843,
            -0.19035768508911133,
            0.015364252030849457,
            -0.483772873878479,
            -0.7650430798530579
        }
    };
    regression_head_bias = {
        0.10385064780712128
    };

    classification_head_weights = {
        {
            -5.264664173126221,
            -1.2807879447937012,
            -0.7477145791053772,
            3.031341552734375,
            1.5341613292694092,
            -1.6744040250778198,
            0.6014524698257446,
            0.12813159823417664,
            1.2626376152038574,
            0.45876526832580566
        },
        {
            -4.001405239105225,
            -1.301662802696228,
            -1.7746647596359253,
            -2.1752195358276367,
            -1.0226296186447144,
            -2.288923740386963,
            -1.088704228401184,
            0.23038308322429657,
            -1.1747989654541016,
            -1.6259526014328003
        },
        {
            -2.5064103603363037,
            2.4615283012390137,
            0.9524018168449402,
            -2.5877861976623535,
            -4.780744552612305,
            3.0858469009399414,
            -5.957152366638184,
            0.1430843323469162,
            0.7261406183242798,
            -0.6905933022499084
        },
        {
            5.6935625076293945,
            0.2643134593963623,
            0.220169335603714,
            -3.896728515625,
            -0.357934832572937,
            0.5103702545166016,
            0.24656452238559723,
            -0.2184055745601654,
            -1.5201313495635986,
            -0.18920253217220306
        }
    };
    classification_head_bias = {
        0.4704197347164154,
        -0.6527161002159119,
        0.4473995566368103,
        0.017734549939632416
    };
}

void sliding_window(void){  
    for (int i = 0; i < input.size(); i++)          //iterate feature
    {
        for (int j = 0; j < input[i].size() - 1; j++) //iterate timestep
        {
        input[i][j] = input[i][j + 1];
        }
        input[i].back() = Voltage_scaled[i];
    }
}

float MultiTaskModel::relu1d(float x) {
    return std::max(0.0f, x);
}

std::vector<std::vector<float>> MultiTaskModel::relu2d(const std::vector<std::vector<float>>& input) {
    std::vector<std::vector<float>> output = input;

    for (auto& channel : output) {      // iterate channel
        for (auto& val : channel) {     // iterate step
            val = std::max(0.0f, val);  // relu
        }
    }

    return output;
}


std::vector<float> MultiTaskModel::softmax(const std::vector<float>& logits) {
    std::vector<float> exp_values;
    float max_logit = *std::max_element(logits.begin(), logits.end());
    float sum_exp = 0.0f;
    for (auto logit : logits) {
        float exp_val = std::exp(logit - max_logit);
        exp_values.push_back(exp_val);
        sum_exp += exp_val;
    }
    for (auto& val : exp_values) {
        val /= sum_exp;
    }
    return exp_values;
}

//outchannel = 1
std::vector<float> MultiTaskModel::conv1d_single_channel(const std::vector<std::vector<float>>& input, 
                                             const std::vector<std::vector<std::vector<float>>>& weights, 
                                             const std::vector<float>& bias, 
                                             int stride, int padding, int kernel_size) {
    //input.shape must be equal to (num_feature, time_step)
    int input_size = input[0].size();
    int num_features = input.size();
    int output_size = (input_size + 2 * padding - kernel_size) / stride + 1;
    
    std::vector<float> output(output_size, 0.0f);

    for (int i = 0; i < output_size; i++) {     //iterate timestep implicitly
        float sum = 0.0f;
        for (int j = 0; j < kernel_size; j++) { //iterate kernel
            int input_index = i * stride + j - padding;
            if (input_index >= 0 && input_index < input_size) {
                for (int input_feature = 0; input_feature < num_features; input_feature++) {    //iterate feature
                    sum += input[input_feature][input_index] * weights[0][input_feature][j];
                }
            }
        }
        output[i] = sum + bias[0];
    }
    return output;
}

std::vector<std::vector<float>> MultiTaskModel::conv1d(const std::vector<std::vector<float>>& input, 
                                                       const std::vector<std::vector<std::vector<float>>>& weights, 
                                                       const std::vector<float>& bias, 
                                                       int stride, int padding, int kernel_size) {
    // input.shape must be (num_features, time_steps)
    int input_size = input[0].size();
    int num_features = input.size();
    int num_output_channels = weights.size(); // number of output channels
    int output_size = (input_size + 2 * padding - kernel_size) / stride + 1;
    
    std::vector<std::vector<float>> output(num_output_channels, std::vector<float>(output_size, 0.0f));

    for (int out_channel = 0; out_channel < num_output_channels; out_channel++) {
        for (int i = 0; i < output_size; i++) {         //iterate timestep implicitly
            float sum = 0.0f;
            for (int j = 0; j < kernel_size; j++) {     //iterate kernel
                int input_index = i * stride + j - padding;
                if (input_index >= 0 && input_index < input_size) {
                    for (int input_feature = 0; input_feature < num_features; input_feature++) {    //iterate feature
                        sum += input[input_feature][input_index] * weights[out_channel][input_feature][j];
                    }
                }
            }
            output[out_channel][i] = sum + bias[out_channel];
        }
    }
    return output;
}

std::vector<float> MultiTaskModel::max_pool1d(const std::vector<float>& input, int kernel_size) {
    int output_size = input.size() / kernel_size;
    std::vector<float> output(output_size, 0.0f);

    for (int j = 0; j < output_size; j++) {
        float max_val = *std::max_element(input.begin() + j * kernel_size, input.begin() + (j + 1) * kernel_size);
        output[j] = max_val;
    }
    return output;
}

std::vector<std::vector<float>> MultiTaskModel::max_pool2d(const std::vector<std::vector<float>>& input, int kernel_size) {
    int num_channels = input.size();
    int output_size = input[0].size() / kernel_size;
    
    std::vector<std::vector<float>> output(num_channels, std::vector<float>(output_size, 0.0f));

    for (int channel = 0; channel < num_channels; channel++) {
        for (int j = 0; j < output_size; j++) {
            float max_val = *std::max_element(input[channel].begin() + j * kernel_size, input[channel].begin() + (j + 1) * kernel_size);
            output[channel][j] = max_val;
        }
    }
    return output;
}

std::vector<float> MultiTaskModel::flatten(const std::vector<std::vector<float>>& input) {
    int num_channels = input.size();
    int flattened_size = 0;
    for (const auto& channel : input) {
        flattened_size += channel.size();
    }

    std::vector<float> flattened_input(flattened_size, 0.0f);

    // Flatten the input
    int index = 0;
    for (int channel = 0; channel < num_channels; channel++) {
        for (int i = 0; i < input[channel].size(); i++) {
            flattened_input[index++] = input[channel][i];
        }
    }

    return flattened_input;
}


// input dimension = 1
std::vector<float> MultiTaskModel::fully_connected(const std::vector<float>& input, 
                                                     const std::vector<std::vector<float>>& weights, 
                                                     const std::vector<float>& bias) {

    std::vector<float> output(bias.size(), 0.0f);

    for (size_t i = 0; i < weights.size(); i++) {
        for (size_t j = 0; j < input.size(); j++) {
            output[i] += input[j] * weights[i][j];
        }
        output[i] += bias[i];
    }
    return output;
}

std::pair<float, std::vector<float>> MultiTaskModel::forward(const std::vector<std::vector<float>>& input) {
    auto conv_output = conv1d(input, conv1_weights, conv1_bias, 1, 1, 3);
    
    conv_output = relu2d(conv_output);

    auto pooled_output = max_pool2d(conv_output, 2);

    // Flatten the pooled output
    auto flattened_output = flatten(pooled_output);

    auto shared_output = fully_connected(flattened_output, shared_weights, shared_bias);

    for (auto& val : shared_output) {
        val = relu1d(val);
    }

    auto regression_output = fully_connected(shared_output, regression_head_weights, regression_head_bias);

    auto classification_output = fully_connected(shared_output, classification_head_weights, classification_head_bias);

    classification_output = softmax(classification_output);

    return {regression_output[0], classification_output}; 
}

void model_forward(void){
	std::pair<float, std::vector<float>> outputs = model.forward(input);
  regression_output = outputs.first;
  std::vector<float> classification_output = outputs.second;
  regression_output = regression_output * y_scaler_scale + y_scaler_mean;

  int max_index = 0;
  float max_value = classification_output[0];
  for (int i = 1; i < classification_output.size(); i++)
  {
      if (classification_output[i] > max_value)
      {
          max_value = classification_output[i];
          max_index = i;
      }
  }
  switch (max_index)
  {
      case 0:
          label = "Air";
          break;
      case 1:
          label = "CO";
          break;
      case 2:
          label = "Ea";
          break;
      case 3:
          label = "H2";
          break;
      default:
          label = "Unknown";
          break;
  }

  if (regression_output <= 0 || label == "Air")
  {
      regression_output = 0;
      label = "Air";
  }
}

